package com.lee.lure.server.ai.aiclient.milvus;

import io.milvus.client.MilvusClient;
import io.milvus.client.MilvusServiceClient;
import io.milvus.param.ConnectParam;

import java.util.concurrent.BlockingQueue;
import java.util.concurrent.LinkedBlockingQueue;

/**
 * @author yanwei li create  2024-07-03 20:20
 */
public class MilvusConnectionPool {

    private final BlockingQueue<MilvusClient> pool;
    private final int poolSize;
    private final String host;
    private final int port;



    public MilvusConnectionPool(int poolSize,String host, int port) {
        this.poolSize = poolSize;
        this.pool = new LinkedBlockingQueue<>(poolSize);
        this.host = host;
        this.port = port;
        for (int i = 0; i < poolSize; i++) {
            MilvusClient client = createClient();
            pool.offer(client);
        }
    }

    private MilvusClient createClient(){
        ConnectParam connectParam = ConnectParam.newBuilder()
                .withHost(host)
                .withPort(port).build();
        return new MilvusServiceClient(connectParam);
    }

    public MilvusClient getConnection() throws InterruptedException {
        return pool.take();
    }

    public void releaseConnection(MilvusClient client) {
        if (client != null) {
            pool.offer(client);
        }
    }

    public void shutdown() {
        for (MilvusClient client : pool) {
            client.close();
        }
    }
}
